#include <Columns/ColumnArray.h>
#include <Columns/ColumnString.h>
#include <Columns/ColumnVector.h>
#include <DataTypes/DataTypeArray.h>
#include <DataTypes/DataTypeString.h>
#include <DataTypes/DataTypesNumber.h>
#include <Functions/FunctionFactory.h>
#include <Functions/IFunction.h>
#include <IO/WriteBufferFromVector.h>
#include <IO/WriteHelpers.h>
#include <bit>


namespace DB
{

namespace ErrorCodes
{
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
    extern const int ILLEGAL_COLUMN;
}


/** Functions for an unusual conversion to a string or array:
  *
  * bitmaskToList - takes an integer - a bitmask, returns a string of degrees of 2 separated by a comma.
  *                 for example, bitmaskToList(50) = '2,16,32'
  *
  * bitmaskToArray(x) - Returns an array of powers of two in the binary form of x. For example, bitmaskToArray(50) = [2, 16, 32].
  *
  */

namespace
{

class FunctionBitmaskToList : public IFunction
{
public:
    static constexpr auto name = "bitmaskToList";
    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionBitmaskToList>(); }

    String getName() const override
    {
        return name;
    }

    size_t getNumberOfArguments() const override { return 1; }
    bool isInjective(const ColumnsWithTypeAndName &) const override { return true; }
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }

    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
    {
        const DataTypePtr & type = arguments[0];

        if (!isInteger(type))
            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Cannot format {} as bitmask string", type->getName());

        return std::make_shared<DataTypeString>();
    }

    DataTypePtr getReturnTypeForDefaultImplementationForDynamic() const override
    {
        return std::make_shared<DataTypeString>();
    }

    bool useDefaultImplementationForConstants() const override { return true; }

    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
    {
        ColumnPtr res;
        if (!((res = executeType<UInt8>(arguments, input_rows_count))
            || (res = executeType<UInt16>(arguments, input_rows_count))
            || (res = executeType<UInt32>(arguments, input_rows_count))
            || (res = executeType<UInt64>(arguments, input_rows_count))
            || (res = executeType<Int8>(arguments, input_rows_count))
            || (res = executeType<Int16>(arguments, input_rows_count))
            || (res = executeType<Int32>(arguments, input_rows_count))
            || (res = executeType<Int64>(arguments, input_rows_count))))
            throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Illegal column {} of argument of function {}",
                            arguments[0].column->getName(), getName());

        return res;
    }

private:
    template <typename T>
    static void writeBitmask(T x, WriteBuffer & out)
    {
        using UnsignedT = make_unsigned_t<T>;
        UnsignedT u_x = x;

        bool first = true;
        while (u_x)
        {
            UnsignedT y = u_x & (u_x - 1);
            UnsignedT bit = u_x ^ y;
            u_x = y;
            if (!first)
                writeChar(',', out);
            first = false;
            writeIntText(static_cast<T>(bit), out);
        }
    }

    template <typename T>
    ColumnPtr executeType(const ColumnsWithTypeAndName & columns, size_t input_rows_count) const
    {
        if (const ColumnVector<T> * col_from = checkAndGetColumn<ColumnVector<T>>(columns[0].column.get()))
        {
            auto col_to = ColumnString::create();

            const typename ColumnVector<T>::Container & vec_from = col_from->getData();
            ColumnString::Chars & data_to = col_to->getChars();
            ColumnString::Offsets & offsets_to = col_to->getOffsets();
            data_to.resize(input_rows_count * 2);
            offsets_to.resize(input_rows_count);

            WriteBufferFromVector<ColumnString::Chars> buf_to(data_to);

            for (size_t i = 0; i < input_rows_count; ++i)
            {
                writeBitmask<T>(vec_from[i], buf_to);
                offsets_to[i] = buf_to.count();
            }

            buf_to.finalize();
            return col_to;
        }

        return nullptr;
    }
};


class FunctionBitmaskToArray : public IFunction
{
public:
    static constexpr auto name = "bitmaskToArray";
    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionBitmaskToArray>(); }

    String getName() const override
    {
        return name;
    }

    size_t getNumberOfArguments() const override { return 1; }
    bool isInjective(const ColumnsWithTypeAndName &) const override { return true; }
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }

    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
    {
        if (!isInteger(arguments[0]))
            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Illegal type {} of argument of function {}",
                            arguments[0]->getName(), getName());

        return std::make_shared<DataTypeArray>(arguments[0]);
    }

    bool useDefaultImplementationForConstants() const override { return true; }

    template <typename T>
    bool tryExecute(const IColumn * column, ColumnPtr & out_column) const
    {
        using UnsignedT = make_unsigned_t<T>;

        if (const ColumnVector<T> * col_from = checkAndGetColumn<ColumnVector<T>>(column))
        {
            auto col_values = ColumnVector<T>::create();
            auto col_offsets = ColumnArray::ColumnOffsets::create();

            typename ColumnVector<T>::Container & res_values = col_values->getData();
            ColumnArray::Offsets & res_offsets = col_offsets->getData();

            const typename ColumnVector<T>::Container & vec_from = col_from->getData();
            size_t size = vec_from.size();
            res_offsets.resize(size);
            res_values.reserve(size * 2);

            for (size_t row = 0; row < size; ++row)
            {
                UnsignedT x = vec_from[row];
                while (x)
                {
                    UnsignedT y = x & (x - 1);
                    UnsignedT bit = x ^ y;
                    x = y;
                    res_values.push_back(bit);
                }
                res_offsets[row] = res_values.size();
            }

            out_column = ColumnArray::create(std::move(col_values), std::move(col_offsets));
            return true;
        }

        return false;
    }

    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t /*input_rows_count*/) const override
    {
        const IColumn * in_column = arguments[0].column.get();
        ColumnPtr out_column;

        if (tryExecute<UInt8>(in_column, out_column) ||
            tryExecute<UInt16>(in_column, out_column) ||
            tryExecute<UInt32>(in_column, out_column) ||
            tryExecute<UInt64>(in_column, out_column) ||
            tryExecute<Int8>(in_column, out_column) ||
            tryExecute<Int16>(in_column, out_column) ||
            tryExecute<Int32>(in_column, out_column) ||
            tryExecute<Int64>(in_column, out_column))
            return out_column;

        throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Illegal column {} of first argument of function {}",
                        arguments[0].column->getName(), getName());
    }
};

class FunctionBitPositionsToArray : public IFunction
{
public:
    static constexpr auto name = "bitPositionsToArray";
    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionBitPositionsToArray>(); }

    String getName() const override
    {
        return name;
    }

    size_t getNumberOfArguments() const override { return 1; }
    bool isInjective(const ColumnsWithTypeAndName &) const override { return true; }
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }

    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
    {
        if (!isInteger(arguments[0]))
            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
                            "Illegal type {} of argument of function {}",
                            getName(),
                            arguments[0]->getName());

        return std::make_shared<DataTypeArray>(std::make_shared<DataTypeUInt64>());
    }

    DataTypePtr getReturnTypeForDefaultImplementationForDynamic() const override
    {
        return std::make_shared<DataTypeArray>(std::make_shared<DataTypeUInt64>());
    }

    bool useDefaultImplementationForConstants() const override { return true; }

    template <typename T>
    ColumnPtr executeType(const IColumn * column, size_t input_rows_count) const
    {
        const ColumnVector<T> * col_from = checkAndGetColumn<ColumnVector<T>>(column);
        if (!col_from)
            return nullptr;

        auto result_array_values = ColumnVector<UInt64>::create();
        auto result_array_offsets = ColumnArray::ColumnOffsets::create();

        auto & result_array_values_data = result_array_values->getData();
        auto & result_array_offsets_data = result_array_offsets->getData();

        auto & vec_from = col_from->getData();
        result_array_offsets_data.resize(input_rows_count);
        result_array_values_data.reserve(input_rows_count * 2);

        using UnsignedType = make_unsigned_t<T>;

        for (size_t row = 0; row < input_rows_count; ++row)
        {
            UnsignedType x = static_cast<UnsignedType>(vec_from[row]);

            if constexpr (is_big_int_v<UnsignedType>)
            {
                size_t position = 0;

                while (x)
                {
                    if (x & 1)
                        result_array_values_data.push_back(position);

                    x >>= 1;
                    ++position;
                }
            }
            else
            {
                while (x)
                {
                    /// С++20 char8_t is not an unsigned integral type anymore https://godbolt.org/z/Mqcb7qn58
                    /// and thus you cannot use std::countr_zero on it.
                    if constexpr (std::is_same_v<UnsignedType, UInt8>)
                        result_array_values_data.push_back(std::countr_zero(static_cast<unsigned char>(x)));
                    else
                        result_array_values_data.push_back(std::countr_zero(x));
                    x &= (x - 1);
                }
            }

            result_array_offsets_data[row] = result_array_values_data.size();
        }

        auto result_column = ColumnArray::create(std::move(result_array_values), std::move(result_array_offsets));

        return result_column;
    }

    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
    {
        const IColumn * in_column = arguments[0].column.get();
        ColumnPtr result_column;

        if (!((result_column = executeType<UInt8>(in_column, input_rows_count))
              || (result_column = executeType<UInt16>(in_column, input_rows_count))
              || (result_column = executeType<UInt32>(in_column, input_rows_count))
              || (result_column = executeType<UInt32>(in_column, input_rows_count))
              || (result_column = executeType<UInt64>(in_column, input_rows_count))
              || (result_column = executeType<UInt128>(in_column, input_rows_count))
              || (result_column = executeType<UInt256>(in_column, input_rows_count))
              || (result_column = executeType<Int8>(in_column, input_rows_count))
              || (result_column = executeType<Int16>(in_column, input_rows_count))
              || (result_column = executeType<Int32>(in_column, input_rows_count))
              || (result_column = executeType<Int64>(in_column, input_rows_count))
              || (result_column = executeType<Int128>(in_column, input_rows_count))
              || (result_column = executeType<Int256>(in_column, input_rows_count))))
        {
            throw Exception(ErrorCodes::ILLEGAL_COLUMN,
                            "Illegal column {} of first argument of function {}",
                            arguments[0].column->getName(),
                            getName());
        }

        return result_column;
    }
};

}

REGISTER_FUNCTION(BitToArray)
{
    FunctionDocumentation::Description bitPositionsToArray_description = R"(
This function returns the positions (in ascending order) of the 1 bits in the binary representation of an unsigned integer.
Signed input integers are first casted to an unsigned integer.
    )";
    FunctionDocumentation::Syntax bitPositionsToArray_syntax = "bitPositionsToArray(arg)";
    FunctionDocumentation::Arguments bitPositionsToArray_arguments = {
        {"arg", "An integer value.", {"(U)Int*"}}
    };
    FunctionDocumentation::ReturnedValue bitPositionsToArray_returned_value = {"Returns an array with the ascendingly ordered positions of 1 bits in the binary representation of the input.", {"Array(UInt64)"}};
    FunctionDocumentation::Examples bitPositionsToArray_examples =
    {
        {
            "Single bit set",
            "SELECT bitPositionsToArray(toInt8(1)) AS bit_positions",
            R"(
┌─bit_positions─┐
│ [0]           │
└───────────────┘
            )"
        },
        {
            "All bits set",
            "SELECT bitPositionsToArray(toInt8(-1)) AS bit_positions",
            R"(
┌─bit_positions─────────────┐
│ [0, 1, 2, 3, 4, 5, 6, 7]  │
└───────────────────────────┘
            )"
        }
    };
    FunctionDocumentation::IntroducedIn bitPositionsToArray_introduced_in = {21, 7};
    FunctionDocumentation::Category bitPositionsToArray_category = FunctionDocumentation::Category::Encoding;
    FunctionDocumentation bitPositionsToArray_documentation = {bitPositionsToArray_description, bitPositionsToArray_syntax, bitPositionsToArray_arguments, bitPositionsToArray_returned_value, bitPositionsToArray_examples, bitPositionsToArray_introduced_in, bitPositionsToArray_category};

    FunctionDocumentation::Description bitmaskToArray_description = R"(
This function decomposes an integer into a sum of powers of two.
The powers of two are returned as an ascendingly ordered array.
    )";
    FunctionDocumentation::Syntax bitmaskToArray_syntax = "bitmaskToArray(num)";
    FunctionDocumentation::Arguments bitmaskToArray_arguments = {{"num", "An integer value.", {"(U)Int*"}}};
    FunctionDocumentation::ReturnedValue bitmaskToArray_returned_value = {"Returns an array with the ascendingly ordered powers of two which sum up to the input number.", {"Array(UInt64)"}};
    FunctionDocumentation::Examples bitmaskToArray_examples = {
        {
            "Basic example",
            "SELECT bitmaskToArray(50) AS powers_of_two",
            R"(
┌─powers_of_two───┐
│ [2, 16, 32]     │
└─────────────────┘
            )"
        },
        {
            "Single power of two",
            "SELECT bitmaskToArray(8) AS powers_of_two",
            R"(
┌─powers_of_two─┐
│ [8]           │
└───────────────┘
            )"
        },
    };
    FunctionDocumentation::IntroducedIn bitmaskToArray_introduced_in = {1, 1};
    FunctionDocumentation::Category bitmaskToArray_category = FunctionDocumentation::Category::Encoding;
    FunctionDocumentation bitmaskToArray_documentation = {bitmaskToArray_description, bitmaskToArray_syntax, bitmaskToArray_arguments, bitmaskToArray_returned_value, bitmaskToArray_examples, bitmaskToArray_introduced_in, bitmaskToArray_category};

    FunctionDocumentation::Description bitmaskToList_description = R"(
Like bitmaskToArray but returns the powers of two as a comma-separated string.
    )";
    FunctionDocumentation::Syntax bitmaskToList_syntax = "bitmaskToList(num)";
    FunctionDocumentation::Arguments bitmaskToList_arguments = {
        {"num", "An integer value.", {"(U)Int*"}}
    };
    FunctionDocumentation::ReturnedValue bitmaskToList_returned_value = {"Returns a string containing comma-separated powers of two.", {"String"}};
    FunctionDocumentation::Examples bitmaskToList_examples = {
        {
            "Basic example", "SELECT bitmaskToList(50) AS powers_list",
            R"(
┌─powers_list───┐
│ 2, 16, 32     │
└───────────────┘
           )"
        },
    };
    FunctionDocumentation::IntroducedIn bitmaskToList_introduced_in = {1, 1};
    FunctionDocumentation::Category bitmaskToList_category = FunctionDocumentation::Category::Encoding;
    FunctionDocumentation bitmaskToList_documentation = {bitmaskToList_description, bitmaskToList_syntax, bitmaskToList_arguments, bitmaskToList_returned_value, bitmaskToList_examples, bitmaskToList_introduced_in, bitmaskToList_category};

    factory.registerFunction<FunctionBitPositionsToArray>(bitPositionsToArray_documentation);
    factory.registerFunction<FunctionBitmaskToArray>(bitmaskToArray_documentation);
    factory.registerFunction<FunctionBitmaskToList>(bitmaskToList_documentation);
}

}
